// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//

#pragma once

#include <algorithm>  // IWYU pragma: keep
#include <cstdint>
#include <limits>
#include <memory>
#include <utility>
#include <vector>

#include <nebula/array/array_base.h>
#include <nebula/array/array_primitive.h>
#include <nebula/core/buffer.h>
#include <nebula/core/buffer_builder.h>

#include <turbo/utility/status.h>
#include <nebula/types/type_fwd.h>
#include <turbo/base/macros.h>

namespace nebula {

    namespace internal {

        template<class Builder, class V>
        class ArrayBuilderExtraOps {
        public:
            /// \brief append a value from an optional or null if it has no value.
            turbo::Status AppendOrNull(const std::optional<V> &value) {
                auto *self = static_cast<Builder *>(this);
                return value.has_value() ? self->append(*value) : self->append_null();
            }

            /// \brief append a value from an optional or null if it has no value.
            ///
            /// Unsafe methods don't check existing size.
            void UnsafeAppendOrNull(const std::optional<V> &value) {
                auto *self = static_cast<Builder *>(this);
                return value.has_value() ? self->unsafe_append(*value) : self->unsafe_append_null();
            }
        };

    }  // namespace internal

    /// \defgroup numeric-builders Concrete builder subclasses for numeric types
    /// @{
    /// @}

    /// \defgroup temporal-builders Concrete builder subclasses for temporal types
    /// @{
    /// @}

    /// \defgroup binary-builders Concrete builder subclasses for binary types
    /// @{
    /// @}

    /// \defgroup nested-builders Concrete builder subclasses for nested types
    /// @{
    /// @}

    /// \defgroup dictionary-builders Concrete builder subclasses for dictionary types
    /// @{
    /// @}

    /// \defgroup run-end-encoded-builders Concrete builder subclasses for run-end encoded
    /// arrays
    /// @{
    /// @}

    constexpr int64_t kMinBuilderCapacity = 1 << 5;
    constexpr int64_t kListMaximumElements = std::numeric_limits<int32_t>::max() - 1;

    /// Base class for all data array builders.
    ///
    /// This class provides a facilities for incrementally building the null bitmap
    /// (see append methods) and as a side effect the current number of slots and
    /// the null count.
    ///
    /// \note Users are expected to use builders as one of the concrete types below.
    /// For example, ArrayBuilder* pointing to BinaryBuilder should be downcast before use.
    class TURBO_EXPORT ArrayBuilder {
    public:
        explicit ArrayBuilder(MemoryPool *pool, int64_t alignment = kDefaultBufferAlignment)
                : pool_(pool), alignment_(alignment), null_bitmap_builder_(pool, alignment) {}

        TURBO_DEFAULT_MOVE_AND_ASSIGN(ArrayBuilder);

        virtual ~ArrayBuilder() = default;

        /// For nested types. Since the objects are owned by this class instance, we
        /// skip shared pointers and just return a raw pointer
        ArrayBuilder *child(int i) { return children_[i].get(); }

        const std::shared_ptr<ArrayBuilder> &child_builder(int i) const { return children_[i]; }

        int num_children() const { return static_cast<int>(children_.size()); }

        virtual int64_t length() const { return length_; }

        int64_t null_count() const { return null_count_; }

        int64_t capacity() const { return capacity_; }

        /// \brief Ensure that enough memory has been allocated to fit the indicated
        /// number of total elements in the builder, including any that have already
        /// been appended. Does not account for reallocations that may be due to
        /// variable size data, like binary values. To make space for incremental
        /// appends, use Reserve instead.
        ///
        /// \param[in] capacity the minimum number of total array values to
        ///            accommodate. Must be greater than the current capacity.
        /// \return turbo::Status
        virtual turbo::Status resize(int64_t capacity);

        /// \brief Ensure that there is enough space allocated to append the indicated
        /// number of elements without any further reallocation. Overallocation is
        /// used in order to minimize the impact of incremental Reserve() calls.
        /// Note that additional_capacity is relative to the current number of elements
        /// rather than to the current capacity, so calls to Reserve() which are not
        /// interspersed with addition of new elements may not increase the capacity.
        ///
        /// \param[in] additional_capacity the number of additional array values
        /// \return turbo::Status
        turbo::Status Reserve(int64_t additional_capacity) {
            auto current_capacity = capacity();
            auto min_capacity = length() + additional_capacity;
            if (min_capacity <= current_capacity) return turbo::OkStatus();

            // leave growth factor up to BufferBuilder
            auto new_capacity = BufferBuilder::GrowByFactor(current_capacity, min_capacity);
            return resize(new_capacity);
        }

        /// reset the builder.
        virtual void reset();

        /// \brief append a null value to builder
        virtual turbo::Status append_null() = 0;

        /// \brief append a number of null values to builder
        virtual turbo::Status append_nulls(int64_t length) = 0;

        /// \brief append a non-null value to builder
        ///
        /// The appended value is an implementation detail, but the corresponding
        /// memory slot is guaranteed to be initialized.
        /// This method is useful when appending a null value to a parent nested type.
        virtual turbo::Status append_empty_value() = 0;

        /// \brief append a number of non-null values to builder
        ///
        /// The appended values are an implementation detail, but the corresponding
        /// memory slot is guaranteed to be initialized.
        /// This method is useful when appending null values to a parent nested type.
        virtual turbo::Status append_empty_values(int64_t length) = 0;

        /// \brief append a value from a scalar
        turbo::Status append_scalar(const Scalar &scalar) { return append_scalar(scalar, 1); }

        virtual turbo::Status append_scalar(const Scalar &scalar, int64_t n_repeats);

        virtual turbo::Status append_scalars(const ScalarVector &scalars);

        /// \brief append a range of values from an array.
        ///
        /// The given array must be the same type as the builder.
        virtual turbo::Status append_array_slice(const ArraySpan & TURBO_ARG_UNUSED(array),
                                               int64_t TURBO_ARG_UNUSED(offset),
                                               int64_t TURBO_ARG_UNUSED(length)) {
            return turbo::unimplemented_error("append_array_slice for builder for ", *type());
        }

        /// \brief Return result of builder as an internal generic ArrayData
        /// object. Resets builder except for dictionary builder
        ///
        /// \param[out] out the finalized ArrayData object
        /// \return turbo::Status
        virtual turbo::Status finish_internal(std::shared_ptr<ArrayData> *out) = 0;

        /// \brief Return result of builder as an Array object.
        ///
        /// The builder is reset except for DictionaryBuilder.
        ///
        /// \param[out] out the finalized Array object
        /// \return turbo::Status
        turbo::Status finish(std::shared_ptr<Array> *out);

        /// \brief Return result of builder as an Array object.
        ///
        /// The builder is reset except for DictionaryBuilder.
        ///
        /// \return The finalized Array object
        turbo::Result<std::shared_ptr<Array>> finish();

        /// \brief Return the type of the built Array
        virtual std::shared_ptr<DataType> type() const = 0;

    protected:
        /// append to null bitmap
        turbo::Status append_to_bitmap(bool is_valid);

        /// Vector append. Treat each zero byte as a null.   If valid_bytes is null
        /// assume all of length bits are valid.
        turbo::Status append_to_bitmap(const uint8_t *valid_bytes, int64_t length);

        /// Uniform append.  append N times the same validity bit.
        turbo::Status append_to_bitmap(int64_t num_bits, bool value);

        /// Set the next length bits to not null (i.e. valid).
        turbo::Status SetNotNull(int64_t length);

        // Unsafe operations (don't check capacity/don't resize)

        void unsafe_append_null() { unsafe_append_to_bitmap(false); }

        // append to null bitmap, update the length
        void unsafe_append_to_bitmap(bool is_valid) {
            null_bitmap_builder_.unsafe_append(is_valid);
            ++length_;
            if (!is_valid) ++null_count_;
        }

        // Vector append. Treat each zero byte as a nullzero. If valid_bytes is null
        // assume all of length bits are valid.
        void unsafe_append_to_bitmap(const uint8_t *valid_bytes, int64_t length) {
            if (valid_bytes == nullptr) {
                return unsafe_set_not_null(length);
            }
            null_bitmap_builder_.unsafe_append(valid_bytes, length);
            length_ += length;
            null_count_ = null_bitmap_builder_.false_count();
        }

        // Vector append. Copy from a given bitmap. If bitmap is null assume
        // all of length bits are valid.
        void unsafe_append_to_bitmap(const uint8_t *bitmap, int64_t offset, int64_t length) {
            if (bitmap == nullptr) {
                return unsafe_set_not_null(length);
            }
            null_bitmap_builder_.unsafe_append(bitmap, offset, length);
            length_ += length;
            null_count_ = null_bitmap_builder_.false_count();
        }

        // append the same validity value a given number of times.
        void unsafe_append_to_bitmap(const int64_t num_bits, bool value) {
            if (value) {
                unsafe_set_not_null(num_bits);
            } else {
                unsafe_set_null(num_bits);
            }
        }

        void unsafe_append_to_bitmap(const std::vector<bool> &is_valid);

        // Set the next validity bits to not null (i.e. valid).
        void unsafe_set_not_null(int64_t length);

        // Set the next validity bits to null (i.e. invalid).
        void unsafe_set_null(int64_t length);

        static turbo::Status trim_buffer(const int64_t bytes_filled, ResizableBuffer *buffer);

        /// \brief finish to an array of the specified ArrayType
        template<typename ArrayType>
        turbo::Status finish_typed(std::shared_ptr<ArrayType> *out) {
            std::shared_ptr<Array> out_untyped;
            TURBO_RETURN_NOT_OK(finish(&out_untyped));
            *out = std::static_pointer_cast<ArrayType>(std::move(out_untyped));
            return turbo::OkStatus();
        }

        // Check the requested capacity for validity
        turbo::Status check_capacity(int64_t new_capacity) {
            if (TURBO_UNLIKELY(new_capacity < 0)) {
                return turbo::invalid_argument_error(
                        "resize capacity must be positive (requested: ", new_capacity, ")");
            }

            if (TURBO_UNLIKELY(new_capacity < length_)) {
                return turbo::invalid_argument_error("resize cannot downsize (requested: ", new_capacity,
                                                     ", current length: ", length_, ")");
            }

            return turbo::OkStatus();
        }

        // Check for array type
        turbo::Status check_array_type(const std::shared_ptr<DataType> &expected_type,
                                     const Array &array, const char *message);

        turbo::Status check_array_type(Type::type expected_type, const Array &array,
                                     const char *message);

        MemoryPool *pool_;
        int64_t alignment_;

        TypedBufferBuilder<bool> null_bitmap_builder_;
        int64_t null_count_ = 0;

        // Array length, so far. Also, the index of the next element to be added
        int64_t length_ = 0;
        int64_t capacity_ = 0;

        // Child value array builders. These are owned by this class
        std::vector<std::shared_ptr<ArrayBuilder>> children_;

    private:
        TURBO_DISALLOW_COPY_AND_ASSIGN(ArrayBuilder);
    };

    /// \brief Construct an empty ArrayBuilder corresponding to the data
    /// type
    /// \param[in] pool the MemoryPool to use for allocations
    /// \param[in] type the data type to create the builder for
    /// \param[out] out the created ArrayBuilder
    TURBO_EXPORT
    turbo::Status MakeBuilder(MemoryPool *pool, const std::shared_ptr<DataType> &type,
                              std::unique_ptr<ArrayBuilder> *out);

    inline turbo::Result<std::unique_ptr<ArrayBuilder>> MakeBuilder(
            const std::shared_ptr<DataType> &type, MemoryPool *pool = default_memory_pool()) {
        std::unique_ptr<ArrayBuilder> out;
        TURBO_RETURN_NOT_OK(MakeBuilder(pool, type, &out));
        return out;
    }

    /// \brief Construct an empty ArrayBuilder corresponding to the data
    /// type, where any top-level or nested dictionary builders return the
    /// exact index type specified by the type.
    TURBO_EXPORT
    turbo::Status MakeBuilderExactIndex(MemoryPool *pool, const std::shared_ptr<DataType> &type,
                                        std::unique_ptr<ArrayBuilder> *out);

    inline turbo::Result<std::unique_ptr<ArrayBuilder>> MakeBuilderExactIndex(
            const std::shared_ptr<DataType> &type, MemoryPool *pool = default_memory_pool()) {
        std::unique_ptr<ArrayBuilder> out;
        TURBO_RETURN_NOT_OK(MakeBuilderExactIndex(pool, type, &out));
        return out;
    }

    /// \brief Construct an empty DictionaryBuilder initialized optionally
    /// with a preexisting dictionary
    /// \param[in] pool the MemoryPool to use for allocations
    /// \param[in] type the dictionary type to create the builder for
    /// \param[in] dictionary the initial dictionary, if any. May be nullptr
    /// \param[out] out the created ArrayBuilder
    TURBO_EXPORT
    turbo::Status MakeDictionaryBuilder(MemoryPool *pool, const std::shared_ptr<DataType> &type,
                                        const std::shared_ptr<Array> &dictionary,
                                        std::unique_ptr<ArrayBuilder> *out);

    inline turbo::Result<std::unique_ptr<ArrayBuilder>> MakeDictionaryBuilder(
            const std::shared_ptr<DataType> &type, const std::shared_ptr<Array> &dictionary,
            MemoryPool *pool = default_memory_pool()) {
        std::unique_ptr<ArrayBuilder> out;
        TURBO_RETURN_NOT_OK(MakeDictionaryBuilder(pool, type, dictionary, &out));
        return out;
    }

}  // namespace nebula
